import os
import json
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

class EncryptionUtils:
    def __init__(self, config_file="keys_config.json"):
        self.config_file = config_file
        self.master_key = self._generate_master_key()
        self.cipher = Fernet(self.master_key)
    
    def _generate_master_key(self):
        """生成主密钥用于加密/解密"""
        # 使用固定的盐值，确保每次启动应用时密钥一致
        salt = b'secret_salt_for_key_management_12345'
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(b"master_password"))
        return key
    
    def encrypt_key(self, key):
        """加密密钥"""
        encrypted_key = self.cipher.encrypt(key.encode())
        return base64.urlsafe_b64encode(encrypted_key).decode()
    
    def decrypt_key(self, encrypted_key):
        """解密密钥"""
        encrypted_data = base64.urlsafe_b64decode(encrypted_key.encode())
        decrypted_key = self.cipher.decrypt(encrypted_data)
        return decrypted_key.decode()
    
    def save_keys(self, keys):
        """保存密钥到配置文件"""
        encrypted_keys = {}
        for name, key in keys.items():
            encrypted_keys[name] = self.encrypt_key(key)
        
        with open(self.config_file, 'w', encoding='utf-8') as f:
            json.dump(encrypted_keys, f, ensure_ascii=False, indent=2)
    
    def load_keys(self):
        """从配置文件加载密钥"""
        if not os.path.exists(self.config_file):
            return {}
        
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                encrypted_keys = json.load(f)
            
            keys = {}
            for name, encrypted_key in encrypted_keys.items():
                try:
                    keys[name] = self.decrypt_key(encrypted_key)
                except Exception:
                    # 如果解密失败，保留原始加密数据
                    keys[name] = encrypted_key
            
            return keys
        except Exception:
            return {}
    
    def key_exists(self):
        """检查配置文件中是否存在密钥"""
        keys = self.load_keys()
        return len(keys) > 0